import gradio as gr
import os
import json
import time
import shutil
import uuid
from typing import Tuple, Optional

from Core.comfy_client import ComfyUIClient
from Core.baidu_translate import baidu_translate


class ReduxLoRAInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "Flux_Redux_Lora.json"
        )

        # 目录：unet、clip、loras、input、temp
        from config.config import get_unet_dir_path, get_clip_dir_path, get_lora_dir_path, get_input_path
        self.unet_dir = get_unet_dir_path()
        self.clip_dir = get_clip_dir_path()
        self.lora_dir = get_lora_dir_path()
        self.input_dir = get_input_path()
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)
        os.makedirs(self.input_dir, exist_ok=True)

        # 模型列表
        self.unet_files = [f for f in os.listdir(self.unet_dir)] if os.path.exists(self.unet_dir) else []
        self.unet_files = [f for f in self.unet_files if f.lower().endswith((".safetensors", ".ckpt"))]
        self.clip_files = [f for f in os.listdir(self.clip_dir)] if os.path.exists(self.clip_dir) else []
        self.clip_files = [f for f in self.clip_files if f.lower().endswith((".safetensors", ".ckpt"))]
        self.lora_choices = self._collect_lora_choices()

        # 默认值
        self.default_unet = "flux1-dev-fp8.safetensors"
        if self.unet_files and self.default_unet not in self.unet_files:
            self.default_unet = self.unet_files[0]
        self.default_clip2 = "t5xxl_fp8_e4m3fn.safetensors"
        if self.clip_files and self.default_clip2 not in self.clip_files:
            self.default_clip2 = self.clip_files[0]
        self.default_lora = "test\\ETShoeF.safetensors"
        if self.lora_choices and self.default_lora not in self.lora_choices:
            self.default_lora = self.lora_choices[0]

    def _collect_lora_choices(self):
        choices = []
        try:
            if os.path.exists(self.lora_dir):
                for root, _, files in os.walk(self.lora_dir):
                    for f in files:
                        if f.lower().endswith((".safetensors", ".ckpt")):
                            rel = os.path.relpath(os.path.join(root, f), self.lora_dir)
                            rel = rel.replace("/", "\\")
                            choices.append(rel)
        except Exception as e:
            print(f"收集LoRA失败: {e}")
        choices.sort()
        return choices

    def load_workflow(self):
        try:
            with open(self.workflow_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"加载工作流失败: {e}")
            return None

    def _save_masked_image(self, base_image_path, mask_data) -> Tuple[Optional[str], Optional[str]]:
        """将蒙版合成到原图的Alpha通道，保存为PNG到input目录并返回文件名与错误信息。
        保障：即使未提供蒙版或蒙版为空，也会生成满alpha(255)的PNG，避免下游节点因空蒙版崩溃。
        """
        try:
            if not base_image_path or not os.path.exists(base_image_path):
                return None, "错误：未提供有效的基础图像路径"
            from PIL import Image
            import numpy as np

            base = Image.open(base_image_path).convert("RGBA")

            # 提取蒙版数组（尽力支持多种结构：dict(mask|layers)、numpy、RGBA的alpha）
            alpha_arr = None
            try:
                if isinstance(mask_data, dict):
                    if "mask" in mask_data and isinstance(mask_data["mask"], np.ndarray):
                        m = mask_data["mask"]
                        # 如果是三通道或四通道，取alpha或转换为灰度
                        if m.ndim == 3:
                            if m.shape[2] == 4:
                                alpha_arr = m[:, :, 3]
                            else:
                                # 取最大通道作为蒙版
                                alpha_arr = m.max(axis=2)
                        elif m.ndim == 2:
                            alpha_arr = m
                    elif "layers" in mask_data and len(mask_data["layers"]) > 0:
                        m = mask_data["layers"][0]
                        if isinstance(m, np.ndarray):
                            if m.ndim == 3:
                                alpha_arr = m[:, :, -1] if m.shape[2] >= 4 else m.max(axis=2)
                            elif m.ndim == 2:
                                alpha_arr = m
                elif isinstance(mask_data, np.ndarray):
                    m = mask_data
                    if m.ndim == 3:
                        if m.shape[2] == 4:
                            alpha_arr = m[:, :, 3]
                        else:
                            alpha_arr = m.max(axis=2)
                    elif m.ndim == 2:
                        alpha_arr = m
            except Exception:
                alpha_arr = None

            # 如果未能提取有效蒙版或蒙版全为0，则生成全255的alpha
            if alpha_arr is None:
                alpha_arr = np.full((base.size[1], base.size[0]), 255, dtype=np.uint8)
            else:
                # 归一化到0/255，并放缩到与原图一致大小
                if alpha_arr.dtype != np.uint8:
                    # 尝试按0-1或其他范围归一化
                    try:
                        alpha_arr = (alpha_arr.astype(np.float32))
                        mmin, mmax = alpha_arr.min(), alpha_arr.max()
                        if mmax > mmin:
                            alpha_arr = (alpha_arr - mmin) / (mmax - mmin)
                        alpha_arr = (alpha_arr * 255.0).clip(0, 255).astype(np.uint8)
                    except Exception:
                        alpha_arr = np.full((base.size[1], base.size[0]), 255, dtype=np.uint8)
                # 如果蒙版非空但全0，避免下游崩溃，改为全255
                if alpha_arr.sum() == 0:
                    alpha_arr = np.full((base.size[1], base.size[0]), 255, dtype=np.uint8)

            # 缩放到与原图相同尺寸
            mask_img = Image.fromarray(alpha_arr)
            if mask_img.size != base.size:
                mask_img = mask_img.resize(base.size)

            r, g, b, _ = base.split()
            out = Image.merge("RGBA", (r, g, b, mask_img))

            filename = f"redux_mask_{uuid.uuid4().hex}.png"
            dst = os.path.join(self.input_dir, filename)
            out.save(dst)
            return filename, None
        except Exception as e:
            return None, f"合成蒙版失败：{e}"

    def _copy_image_to_input(self, src_path, prefix):
        try:
            if not src_path or not os.path.exists(src_path):
                return None, "错误：未提供有效的图像路径"
            ext = os.path.splitext(src_path)[1].lower()
            if ext not in [".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"]:
                return None, "错误：仅支持PNG/JPG/JPEG/BMP/WEBP/TIFF"
            filename = f"{prefix}_{uuid.uuid4().hex}{ext}"
            dst = os.path.join(self.input_dir, filename)
            shutil.copy2(src_path, dst)
            return filename, None
        except Exception as e:
            return None, f"复制图片失败：{e}"

    def generate_image(
        self,
        input_image_path,
        mask_data,
        redux_weight,
        unet_name,
        clip_name2,
        lora_name,
        strength_model,
        width,
        height,
        steps,
        denoise,
        model_ref_image_path,
        text_input,
        max_retries: int = 40,
    ):
        try:
            workflow = self.load_workflow()
            if not workflow:
                return None, "无法加载工作流", None

            # 文本翻译（节点18 TextInput_）
            try:
                english = baidu_translate(text_input or "")
                translate_info = "翻译成功"
            except Exception as e:
                english = text_input or ""
                translate_info = f"翻译失败，改用原文本: {str(e)}"
            if "18" in workflow and "inputs" in workflow["18"]:
                workflow["18"]["inputs"]["text"] = english

            # 节点29：UNET模型
            if "29" in workflow and "inputs" in workflow["29"]:
                workflow["29"]["inputs"]["unet_name"] = unet_name

            # 节点27：DualCLIP clip_name2
            if "27" in workflow and "inputs" in workflow["27"]:
                workflow["27"]["inputs"]["clip_name2"] = clip_name2

            # 节点32：LoRA加载器
            if "32" in workflow and "inputs" in workflow["32"]:
                workflow["32"]["inputs"]["lora_name"] = lora_name
                workflow["32"]["inputs"]["strength_model"] = float(strength_model)

            # 节点34：空Latent尺寸
            if "34" in workflow and "inputs" in workflow["34"]:
                workflow["34"]["inputs"]["width"] = int(width)
                workflow["34"]["inputs"]["height"] = int(height)

            # 节点33：KSampler步数与降噪
            if "33" in workflow and "inputs" in workflow["33"]:
                workflow["33"]["inputs"]["steps"] = int(steps)
                workflow["33"]["inputs"]["denoise"] = float(denoise)

            # 节点20：ReduxAdvanced 权重
            if "20" in workflow and "inputs" in workflow["20"]:
                workflow["20"]["inputs"]["weight"] = float(redux_weight)

            # 节点9：合成蒙版并设置图像（确保为PNG且alpha非空）
            masked_filename, err = self._save_masked_image(input_image_path, mask_data)
            if err:
                return None, f"生成失败：{err}", english
            if "9" in workflow and "inputs" in workflow["9"]:
                workflow["9"]["inputs"]["image"] = masked_filename

            # 如果无法获取有效蒙版，避免ReduxAdvanced在"autocrop with mask"模式下崩溃
            try:
                if "20" in workflow and "inputs" in workflow["20"]:
                    mode = workflow["20"]["inputs"].get("mode", "autocrop with mask")
                    # 这里我们已确保alpha非空，但如果用户未绘制蒙版，仍可切换为无蒙版自适应裁剪
                    if mask_data is None and mode == "autocrop with mask":
                        workflow["20"]["inputs"]["mode"] = "autocrop"
            except Exception:
                pass

            # 节点14：加载模型参考图
            ref_filename, err2 = self._copy_image_to_input(model_ref_image_path, "redux_ref")
            if err2:
                return None, f"生成失败：{err2}", english
            if "14" in workflow and "inputs" in workflow["14"]:
                workflow["14"]["inputs"]["image"] = ref_filename

            # 提交生成请求
            response = self.client.post_prompt(workflow)
            prompt_id = response.get("prompt_id")
            if not prompt_id:
                return None, "生成失败：未返回prompt_id", english

            # 轮询获取结果图
            retry_interval = 2
            for _ in range(max_retries):
                history = self.client.get_history(prompt_id)
                if prompt_id in history and "outputs" in history[prompt_id]:
                    outputs = history[prompt_id]["outputs"]
                    for node_id, data in outputs.items():
                        if "images" in data:
                            for image in data["images"]:
                                img_bytes = self.client.get_image(image["filename"], image.get("subfolder"), image.get("type"))
                                result_file = os.path.join(self.temp_dir, f"final_{prompt_id}.png")
                                with open(result_file, "wb") as f:
                                    f.write(img_bytes)
                                # 同步保存到项目 output_images
                                try:
                                    from config.config import get_local_output_images_path
                                    out_dir = get_local_output_images_path()
                                    os.makedirs(out_dir, exist_ok=True)
                                    shutil.copy2(result_file, os.path.join(out_dir, os.path.basename(result_file)))
                                except Exception as e:
                                    print(f"复制到output_images失败: {e}")
                                status = "生成成功\n" + translate_info
                                return os.path.abspath(result_file), status, english
                time.sleep(retry_interval)

            return None, "生成失败：超时或无输出", english
        except Exception as e:
            print(f"生成异常: {e}")
            return None, f"生成失败：{e}", None


def Tab_redux_LoRA():
    interface = ReduxLoRAInterface()

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 根据LoRA模型和参考图设计")

            # 节点9：加载图片 + 蒙版绘制
            upload_image = gr.Image(label="加载图片", type="filepath", interactive=True)
            image_editor = gr.ImageEditor(label="绘制蒙版", type="numpy", brush=gr.Brush(colors=["#000000"], color_mode="fixed"),interactive=True)

            def update_editor(img_path):
                return img_path

            upload_image.change(fn=update_editor, inputs=upload_image, outputs=image_editor)

            # 节点20：ReduxAdvanced 权重
            redux_weight_slider = gr.Slider(label="Redux权重(weight)", minimum=0, maximum=1, value=1, step=0.01)

            # 节点29：UNET选择
            unet_dropdown = gr.Dropdown(label="UNET模型(unet_name)", choices=interface.unet_files, value=interface.default_unet, interactive=bool(interface.unet_files))

            # 节点27：clip_name2 模型
            clip2_dropdown = gr.Dropdown(label="CLIP文本编码器(clip_name2)", choices=interface.clip_files, value=interface.default_clip2, interactive=bool(interface.clip_files))

            # 节点32：LoRA加载器
            lora_dropdown = gr.Dropdown(label="LoRA名称(lora_name)", choices=interface.lora_choices, value=interface.default_lora, interactive=bool(interface.lora_choices))
            strength_model_slider = gr.Slider(label="LoRA模型强度(strength_model)", minimum=0, maximum=1, value=0.8, step=0.01)

            # 节点34：空latent尺寸
            width_input = gr.Number(label="宽度(width)", value=1024, minimum=0, maximum=2048)
            height_input = gr.Number(label="高度(height)", value=1024, minimum=0, maximum=2048)

            # 节点33：KSampler
            steps_input = gr.Number(label="步数(steps)", value=20, minimum=0, maximum=60)
            denoise_slider = gr.Slider(label="降噪(denoise)", minimum=0, maximum=1, value=1, step=0.01)

            # 节点14：模型参考图
            model_ref_image = gr.Image(label="加载模型参考图", type="filepath")

            # 节点18：文本 + 翻译展示
            text_input = gr.Textbox(label="文本(中文自动翻译)", lines=3)
            translated_text = gr.Textbox(label="英文翻译(用于工作流)", lines=3, interactive=False)

        with gr.Column(scale=2):
            generate_btn = gr.Button("生成", variant="primary")
            result_image = gr.Image(label="生成结果")
            status_text = gr.Textbox(label="状态", interactive=False)

            def get_output_images():
                from config.config import get_local_output_images_path
                out_dir = get_local_output_images_path()
                if not os.path.exists(out_dir):
                    return []
                images = []
                for file in os.listdir(out_dir):
                    if file.lower().endswith((".png", ".jpg", ".jpeg")):
                        images.append(os.path.join(out_dir, file))
                images.sort(key=lambda x: os.path.getmtime(x), reverse=True)
                return images
            
            with gr.Row():
                refresh_btn = gr.Button("刷新输出预览")
            with gr.Row():
                output_gallery = gr.Gallery(label="输出文件夹图片", columns=4, show_label=True, elem_id="output_gallery_i2i")

            refresh_btn.click(fn=get_output_images, outputs=[output_gallery])

    def translate_and_generate(img_path, mask_data, redux_weight, unet_name, clip2_name, lora_name, strength_model, width, height, steps, denoise, model_ref_path, text):
        img_path, status, english = interface.generate_image(
            img_path, mask_data, redux_weight, unet_name, clip2_name, lora_name, strength_model, width, height, steps, denoise, model_ref_path, text
        )
        return img_path, status, english

    generate_btn.click(
        fn=translate_and_generate,
        inputs=[
            upload_image,
            image_editor,
            redux_weight_slider,
            unet_dropdown,
            clip2_dropdown,
            lora_dropdown,
            strength_model_slider,
            width_input,
            height_input,
            steps_input,
            denoise_slider,
            model_ref_image,
            text_input,
        ],
        outputs=[result_image, status_text, translated_text],
    )